package team.bluepen.supermarket.service.calc;

import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.spark.SparkConf;
import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.kafka010.ConsumerStrategies;
import org.apache.spark.streaming.kafka010.KafkaUtils;
import org.apache.spark.streaming.kafka010.LocationStrategies;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import team.bluepen.supermarket.conf.KafkaProperties;
import team.bluepen.supermarket.conf.SparkProperties;
import team.bluepen.supermarket.data.entity.Product;
import team.bluepen.supermarket.service.ProductClientMethodRegister;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

/**
 * @author RollW
 */
public class ProductSparkStreamingClient {
    private static final Logger logger = LoggerFactory.getLogger(ProductSparkStreamingClient.class);

    private final JavaDStream<Product> productDStream;
    private final JavaStreamingContext context;

    public ProductSparkStreamingClient(KafkaProperties kafkaProperties,
                                       SparkProperties sparkProperties,
                                       String appName, String topic,
                                       ProductClientMethodRegister register) {
        SparkConf conf = new SparkConf()
                .setAppName(appName)
                .setMaster(sparkProperties.getMaster());
        context = new JavaStreamingContext(conf, Durations.seconds(4));
        Map<String, Object> kafkaParams = toParams(kafkaProperties);
        JavaInputDStream<ConsumerRecord<String, String>> stream =
                KafkaUtils.createDirectStream(context,
                        LocationStrategies.PreferConsistent(),
                        ConsumerStrategies.Subscribe(
                                Collections.singletonList(topic),
                                kafkaParams)
                );

        productDStream = stream
                .map(record -> Product.deserialize(record.value()));
        register.registerMethods(productDStream);
    }


    public ProductSparkStreamingClient(KafkaProperties kafkaProperties,
                                       SparkContext sparkContext,
                                       String topic,
                                       ProductClientMethodRegister register) {
        context = new JavaStreamingContext(new JavaSparkContext(sparkContext),
                Durations.seconds(4));
        Map<String, Object> kafkaParams = toParams(kafkaProperties);
        JavaInputDStream<ConsumerRecord<String, String>> stream =
                KafkaUtils.createDirectStream(context,
                        LocationStrategies.PreferConsistent(),
                        ConsumerStrategies.Subscribe(
                                Collections.singletonList(topic),
                                kafkaParams)
                );

        productDStream = stream
                .map(record -> Product.deserialize(record.value()));
        register.registerMethods(productDStream);
    }

    private static Map<String, Object> toParams(KafkaProperties properties) {
        Map<String, Object> kafkaParams = new HashMap<>();
        kafkaParams.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, properties.getBootstrapServers());
        kafkaParams.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        kafkaParams.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        kafkaParams.put(ConsumerConfig.GROUP_ID_CONFIG, "SparkClient");
        kafkaParams.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        return kafkaParams;
    }

    public JavaStreamingContext getContext() {
        return context;
    }

    public void start() {
        logger.debug("Start client");
        context.start();
    }
}
